{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "DistributedTripletMarginLossMNIST.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "0SEqLaxOq7kw",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 225
        },
        "outputId": "ebb800a3-3fc8-48c3-bcf5-502e3c5c578c"
      },
      "source": [
        "!pip install pytorch-metric-learning\n",
        "!pip install faiss-cpu # we're using cpu for this example"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: pytorch-metric-learning==0.9.92.dev1 in /usr/local/lib/python3.6/dist-packages (0.9.92.dev1)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from pytorch-metric-learning==0.9.92.dev1) (1.18.5)\n",
            "Requirement already satisfied: scikit-learn in /usr/local/lib/python3.6/dist-packages (from pytorch-metric-learning==0.9.92.dev1) (0.22.2.post1)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.6/dist-packages (from pytorch-metric-learning==0.9.92.dev1) (4.41.1)\n",
            "Requirement already satisfied: torch>=1.6.0 in /usr/local/lib/python3.6/dist-packages (from pytorch-metric-learning==0.9.92.dev1) (1.6.0+cu101)\n",
            "Requirement already satisfied: torchvision in /usr/local/lib/python3.6/dist-packages (from pytorch-metric-learning==0.9.92.dev1) (0.7.0+cu101)\n",
            "Requirement already satisfied: scipy>=0.17.0 in /usr/local/lib/python3.6/dist-packages (from scikit-learn->pytorch-metric-learning==0.9.92.dev1) (1.4.1)\n",
            "Requirement already satisfied: joblib>=0.11 in /usr/local/lib/python3.6/dist-packages (from scikit-learn->pytorch-metric-learning==0.9.92.dev1) (0.16.0)\n",
            "Requirement already satisfied: future in /usr/local/lib/python3.6/dist-packages (from torch>=1.6.0->pytorch-metric-learning==0.9.92.dev1) (0.16.0)\n",
            "Requirement already satisfied: pillow>=4.1.1 in /usr/local/lib/python3.6/dist-packages (from torchvision->pytorch-metric-learning==0.9.92.dev1) (7.0.0)\n",
            "Requirement already satisfied: faiss-cpu in /usr/local/lib/python3.6/dist-packages (1.6.3)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from faiss-cpu) (1.18.5)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vYmfxuw1tYdm",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "outputId": "d7b26d2d-0e00-4418-b210-240c745229b9"
      },
      "source": [
        "######################################################################################################################\n",
        "### This script is modified from the guide on pytorch distributed training https://github.com/seba-1511/dist_tuto.pth/\n",
        "### https://pytorch.org/tutorials/intermediate/dist_tuto.html\n",
        "######################################################################################################################\n",
        "import os\n",
        "import torch\n",
        "import torch.distributed as dist\n",
        "from torch.nn.parallel import DistributedDataParallel as DDP\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "\n",
        "from math import ceil\n",
        "from random import Random\n",
        "from torch.multiprocessing import Process\n",
        "from torchvision import datasets, transforms\n",
        "\n",
        "from pytorch_metric_learning.utils import distributed as pml_dist\n",
        "from pytorch_metric_learning.utils.accuracy_calculator import AccuracyCalculator\n",
        "from pytorch_metric_learning import losses, miners, testers\n",
        "import numpy as np\n",
        "import logging\n",
        "logging.getLogger().setLevel(logging.INFO)\n",
        "\n",
        "\n",
        "class Partition(object):\n",
        "    \"\"\" Dataset-like object, but only access a subset of it. \"\"\"\n",
        "\n",
        "    def __init__(self, data, index):\n",
        "        self.data = data\n",
        "        self.index = index\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.index)\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        data_idx = self.index[index]\n",
        "        return self.data[data_idx]\n",
        "\n",
        "\n",
        "class DataPartitioner(object):\n",
        "    \"\"\" Partitions a dataset into different chuncks. \"\"\"\n",
        "\n",
        "    def __init__(self, data, sizes=[0.7, 0.2, 0.1], seed=1234):\n",
        "        self.data = data\n",
        "        self.partitions = []\n",
        "        rng = Random()\n",
        "        rng.seed(seed)\n",
        "        data_len = len(data)\n",
        "        indexes = [x for x in range(0, data_len)]\n",
        "        rng.shuffle(indexes)\n",
        "\n",
        "        for frac in sizes:\n",
        "            part_len = int(frac * data_len)\n",
        "            self.partitions.append(indexes[0:part_len])\n",
        "            indexes = indexes[part_len:]\n",
        "\n",
        "    def use(self, partition):\n",
        "        return Partition(self.data, self.partitions[partition])\n",
        "\n",
        "\n",
        "class Net(nn.Module):\n",
        "    \"\"\" Network architecture. \"\"\"\n",
        "\n",
        "    def __init__(self):\n",
        "        super(Net, self).__init__()\n",
        "        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)\n",
        "        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)\n",
        "        self.conv2_drop = nn.Dropout2d()\n",
        "        self.fc1 = nn.Linear(320, 50)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = F.relu(F.max_pool2d(self.conv1(x), 2))\n",
        "        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))\n",
        "        x = x.view(-1, 320)\n",
        "        return self.fc1(x)\n",
        "\n",
        "def get_MNIST(train):\n",
        "    return datasets.MNIST(\n",
        "            './data',\n",
        "            train=train,\n",
        "            download=True,\n",
        "            transform=transforms.Compose([\n",
        "                transforms.ToTensor(),\n",
        "                transforms.Normalize((0.1307, ), (0.3081, ))\n",
        "            ]))\n",
        "    \n",
        "\n",
        "def partition_dataset(dataset):\n",
        "    \"\"\" Partitioning MNIST \"\"\"\n",
        "    size = dist.get_world_size()\n",
        "    bsz = 512 // size\n",
        "    partition_sizes = [1.0 / size for _ in range(size)]\n",
        "    partition = DataPartitioner(dataset, partition_sizes)\n",
        "    partition = partition.use(dist.get_rank())\n",
        "    train_set = torch.utils.data.DataLoader(\n",
        "        partition, batch_size=bsz, shuffle=True)\n",
        "    return train_set, bsz\n",
        "\n",
        "\n",
        "### convenient function from pytorch-metric-learning ###\n",
        "def get_all_embeddings(dataset, model, data_device):\n",
        "    # dataloader_num_workers has to be 0 to avoid pid error\n",
        "    # This only happens when within multiprocessing\n",
        "    tester = testers.BaseTester(dataloader_num_workers=0, data_device=data_device)\n",
        "    return tester.get_all_embeddings(dataset, model)\n",
        "\n",
        "\n",
        "### compute accuracy using AccuracyCalculator from pytorch-metric-learning ###\n",
        "def test(dataset, model, accuracy_calculator, data_device):\n",
        "    embeddings, labels = get_all_embeddings(dataset, model, data_device)\n",
        "    print(\"Computing accuracy\")\n",
        "    accuracies = accuracy_calculator.get_accuracy(embeddings, \n",
        "                                                embeddings,\n",
        "                                                np.squeeze(labels),\n",
        "                                                np.squeeze(labels),\n",
        "                                                True)\n",
        "    print(\"Validation set accuracy (MAP@10) = {}\".format(accuracies[\"mean_average_precision_at_r\"]))\n",
        "\n",
        "def test_model(rank, dataset, model, epoch, data_device):\n",
        "    if rank == 0:\n",
        "        print(\"Computing validation set accuracy for epoch {}\".format(epoch))\n",
        "        accuracy_calculator = AccuracyCalculator(include = (\"mean_average_precision_at_r\",), k = 10)\n",
        "        test(dataset, model, accuracy_calculator, data_device)\n",
        "    dist.barrier()\n",
        "\n",
        "\n",
        "def run(rank, size, train_dataset, val_dataset):\n",
        "    \"\"\" Distributed Synchronous SGD Example \"\"\"\n",
        "    print(\"Rank {} entering the 'run' function\".format(rank))\n",
        "    torch.manual_seed(1234)\n",
        "    train_set, bsz = partition_dataset(train_dataset)\n",
        "    dist.barrier()\n",
        "    ### use this if you have multiple GPUs ###\n",
        "    # device = torch.device(\"cuda:{}\".format(rank))\n",
        "    device = torch.device(\"cpu\")\n",
        "    model = Net()\n",
        "    ### if you have multiple GPUs, set this to DDP(model.to(device), device_ids=[rank])\n",
        "    model = DDP(model.to(device))\n",
        "    test_model(rank, val_dataset, model, \"untrained\", device)\n",
        "\n",
        "    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)\n",
        "\n",
        "    #####################################\n",
        "    ### pytorch-metric-learning stuff ###\n",
        "    loss_fn = losses.TripletMarginLoss()\n",
        "    ### if you have multiple GPUs, set this to pml_dist.DistributedLossWrapper(loss=loss_fn, device_ids=[rank])\n",
        "    loss_fn = pml_dist.DistributedLossWrapper(loss=loss_fn)\n",
        "    miner = miners.MultiSimilarityMiner()\n",
        "    miner = pml_dist.DistributedMinerWrapper(miner=miner)\n",
        "    ### pytorch-metric-learning stuff ###\n",
        "    #####################################\n",
        "\n",
        "    num_batches = ceil(len(train_set.dataset) / float(bsz))\n",
        "    for epoch in range(1):\n",
        "        epoch_loss = 0.0\n",
        "        print(\"Rank {} starting epoch {}\".format(rank, epoch))\n",
        "        for i, (data, target) in enumerate(train_set):\n",
        "            data, target = data.to(device), target.to(device)\n",
        "            optimizer.zero_grad()\n",
        "            output = model(data)\n",
        "            hard_pairs = miner(output, target)\n",
        "            loss = loss_fn(output, target, hard_pairs)\n",
        "            epoch_loss += loss.item()\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "            if i % 10 == 0:\n",
        "                print('Rank {}, iteration {}, loss {}, num pos pairs {}, num neg pairs {}'.\\\n",
        "                      format(rank, i, loss.item(), miner.miner.num_pos_pairs, miner.miner.num_neg_pairs))\n",
        "            dist.barrier()\n",
        "\n",
        "        print('Rank {}, epoch {}, average loss {}'.format(rank, epoch, epoch_loss/num_batches))\n",
        "        test_model(rank, val_dataset, model, epoch, device)\n",
        "\n",
        "\n",
        "\n",
        "#######################################\n",
        "### Set backend='nccl' if using GPU ###\n",
        "#######################################\n",
        "def init_processes(rank, size, fn, train_dataset, val_dataset, backend='gloo'):\n",
        "    \"\"\" Initialize the distributed environment. \"\"\"\n",
        "    os.environ['MASTER_ADDR'] = 'localhost'\n",
        "    os.environ['MASTER_PORT'] = '29500'\n",
        "    dist.init_process_group(backend, rank=rank, world_size=size)\n",
        "    fn(rank, size, train_dataset, val_dataset)\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    train_dataset = get_MNIST(True)\n",
        "    val_dataset = get_MNIST(False)\n",
        "\n",
        "    size = 4\n",
        "    processes = []\n",
        "    for rank in range(size):\n",
        "        p = Process(target=init_processes, args=(rank, size, run, train_dataset, val_dataset))\n",
        "        p.start()\n",
        "        processes.append(p)\n",
        "\n",
        "    for p in processes:\n",
        "        p.join()"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Rank 1 entering the 'run' function\n",
            "Rank 2 entering the 'run' function\n",
            "Rank 0 entering the 'run' function\n",
            "Rank 3 entering the 'run' function\n",
            "Computing validation set accuracy for epoch untrained\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 313/313 [00:03<00:00, 88.79it/s]\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Computing accuracy\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:root:running k-nn with k=10\n",
            "INFO:root:embedding dimensionality is 50\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Validation set accuracy (MAP@10) = 0.7912241984126984\n",
            "Rank 1 starting epoch 0\n",
            "Rank 2 starting epoch 0\n",
            "Rank 0 starting epoch 0\n",
            "Rank 3 starting epoch 0\n",
            "Rank 2, iteration 0, loss 0.22585304081439972, num pos pairs 26208, num neg pairs 234298\n",
            "Rank 3, iteration 0, loss 0.22585304081439972, num pos pairs 26208, num neg pairs 234298\n",
            "Rank 0, iteration 0, loss 0.22585304081439972, num pos pairs 26208, num neg pairs 234298\n",
            "Rank 1, iteration 0, loss 0.22585304081439972, num pos pairs 26208, num neg pairs 234298\n",
            "Rank 3, iteration 10, loss 0.20190657675266266, num pos pairs 26239, num neg pairs 233600\n",
            "Rank 2, iteration 10, loss 0.20190657675266266, num pos pairs 26239, num neg pairs 233600\n",
            "Rank 1, iteration 10, loss 0.20190657675266266, num pos pairs 26239, num neg pairs 233600\n",
            "Rank 0, iteration 10, loss 0.20190657675266266, num pos pairs 26239, num neg pairs 233600\n",
            "Rank 3, iteration 20, loss 0.18472252786159515, num pos pairs 26811, num neg pairs 232881\n",
            "Rank 2, iteration 20, loss 0.18472252786159515, num pos pairs 26811, num neg pairs 232881\n",
            "Rank 0, iteration 20, loss 0.18472252786159515, num pos pairs 26811, num neg pairs 232881\n",
            "Rank 1, iteration 20, loss 0.18472252786159515, num pos pairs 26811, num neg pairs 232881\n",
            "Rank 0, iteration 30, loss 0.16923512518405914, num pos pairs 26634, num neg pairs 232605\n",
            "Rank 1, iteration 30, loss 0.16923512518405914, num pos pairs 26634, num neg pairs 232605\n",
            "Rank 2, iteration 30, loss 0.16923512518405914, num pos pairs 26634, num neg pairs 232605\n",
            "Rank 3, iteration 30, loss 0.16923512518405914, num pos pairs 26634, num neg pairs 232605\n",
            "Rank 0, iteration 40, loss 0.16270387172698975, num pos pairs 25952, num neg pairs 233556\n",
            "Rank 2, iteration 40, loss 0.16270387172698975, num pos pairs 25952, num neg pairs 233556\n",
            "Rank 3, iteration 40, loss 0.16270387172698975, num pos pairs 25952, num neg pairs 233556\n",
            "Rank 1, iteration 40, loss 0.16270387172698975, num pos pairs 25952, num neg pairs 233556\n",
            "Rank 3, iteration 50, loss 0.1537935435771942, num pos pairs 25948, num neg pairs 232996\n",
            "Rank 1, iteration 50, loss 0.1537935435771942, num pos pairs 25948, num neg pairs 232996\n",
            "Rank 0, iteration 50, loss 0.1537935435771942, num pos pairs 25948, num neg pairs 232996\n",
            "Rank 2, iteration 50, loss 0.1537935435771942, num pos pairs 25948, num neg pairs 232996\n",
            "Rank 3, iteration 60, loss 0.14659938216209412, num pos pairs 25999, num neg pairs 231592\n",
            "Rank 1, iteration 60, loss 0.14659938216209412, num pos pairs 25999, num neg pairs 231592\n",
            "Rank 0, iteration 60, loss 0.14659938216209412, num pos pairs 25999, num neg pairs 231592\n",
            "Rank 2, iteration 60, loss 0.14659938216209412, num pos pairs 25999, num neg pairs 231592\n",
            "Rank 1, iteration 70, loss 0.14516910910606384, num pos pairs 26212, num neg pairs 231347\n",
            "Rank 0, iteration 70, loss 0.14516910910606384, num pos pairs 26212, num neg pairs 231347\n",
            "Rank 3, iteration 70, loss 0.14516910910606384, num pos pairs 26212, num neg pairs 231347\n",
            "Rank 2, iteration 70, loss 0.14516910910606384, num pos pairs 26212, num neg pairs 231347\n",
            "Rank 0, iteration 80, loss 0.13959552347660065, num pos pairs 25856, num neg pairs 230818\n",
            "Rank 3, iteration 80, loss 0.13959552347660065, num pos pairs 25856, num neg pairs 230818\n",
            "Rank 2, iteration 80, loss 0.13959552347660065, num pos pairs 25856, num neg pairs 230818\n",
            "Rank 1, iteration 80, loss 0.13959552347660065, num pos pairs 25856, num neg pairs 230818\n",
            "Rank 3, iteration 90, loss 0.13859081268310547, num pos pairs 25906, num neg pairs 228747\n",
            "Rank 0, iteration 90, loss 0.13859081268310547, num pos pairs 25906, num neg pairs 228747\n",
            "Rank 1, iteration 90, loss 0.13859081268310547, num pos pairs 25906, num neg pairs 228747\n",
            "Rank 2, iteration 90, loss 0.13859081268310547, num pos pairs 25906, num neg pairs 228747\n",
            "Rank 1, iteration 100, loss 0.13722871243953705, num pos pairs 25611, num neg pairs 228700\n",
            "Rank 2, iteration 100, loss 0.13722871243953705, num pos pairs 25611, num neg pairs 228700\n",
            "Rank 0, iteration 100, loss 0.13722871243953705, num pos pairs 25611, num neg pairs 228700\n",
            "Rank 3, iteration 100, loss 0.13722871243953705, num pos pairs 25611, num neg pairs 228700\n",
            "Rank 0, iteration 110, loss 0.13637162744998932, num pos pairs 26328, num neg pairs 225989\n",
            "Rank 1, iteration 110, loss 0.13637162744998932, num pos pairs 26328, num neg pairs 225989\n",
            "Rank 2, iteration 110, loss 0.13637162744998932, num pos pairs 26328, num neg pairs 225989\n",
            "Rank 3, iteration 110, loss 0.13637162744998932, num pos pairs 26328, num neg pairs 225989\n",
            "Rank 1, epoch 0, average loss 0.15869572645021698\n",
            "Rank 2, epoch 0, average loss 0.15869572645021698\n",
            "Rank 3, epoch 0, average loss 0.15869572645021698\n",
            "Rank 0, epoch 0, average loss 0.15869572645021698\n",
            "Computing validation set accuracy for epoch 0\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 313/313 [00:03<00:00, 89.69it/s]\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Computing accuracy\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:root:running k-nn with k=10\n",
            "INFO:root:embedding dimensionality is 50\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Validation set accuracy (MAP@10) = 0.8194116587301588\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}